from matplotlib import pyplot as plt
from sklearn.model_selection import train_test_split
import tensorflow

import tensorflow as tf
import pandas as pd
import numpy as np

# 读取数据
acceleration_data = pd.read_csv('data_x.csv', header=None)
labels = pd.read_csv('data_y.csv', header=None)

# 将数据转换为 numpy 数组
X = acceleration_data.values
y = labels.values

# 转换形状以适应模型输入
X = X.reshape(-1, 128, 3).astype(np.float32)
y = y.flatten().astype(np.int32)

# 分割数据集为训练集和测试集
from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 将标签进行 one-hot 编码
num_classes = len(np.unique(y))
y_train = tf.keras.utils.to_categorical(y_train, num_classes)
y_test = tf.keras.utils.to_categorical(y_test, num_classes)

# 数据增强函数
# def augment_data(X):
#     augmented_X = []
#     for x in X:
#         augmented_X.append(x)
#         # 添加噪声
#         noise = np.random.normal(0, 0.02, x.shape)
#         augmented_X.append(x + noise)
#         # 翻转
#         augmented_X.append(np.flip(x, axis=0))
#     return np.array(augmented_X)

# # 增强训练数据
# X_train = augment_data(X_train)
# y_train = np.tile(y_train, (3, 1))

# 建立模型
model = tf.keras.Sequential([
    tf.keras.layers.Conv1D(16, 3, activation='relu', input_shape=(128, 3)),
    tf.keras.layers.MaxPooling1D(2),
    tf.keras.layers.Conv1D(32, 3, activation='relu'),
    tf.keras.layers.MaxPooling1D(2),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dropout(0.5),
    tf.keras.layers.Dense(num_classes, activation='softmax')
])

# 编译模型
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

# 打印模型摘要
model.summary()

# 训练模型
history = model.fit(X_train, y_train, epochs=100, batch_size=3, validation_split=0.1)

# 评估模型
loss, accuracy = model.evaluate(X_test, y_test)
print ("Model loss is",loss)
print ("Model accuracy is",accuracy)

converter = tensorflow.lite.TFLiteConverter.from_keras_model(model)
# converter.optimizations = [tensorflow.lite.Optimize.DEFAULT]
# converter.target_spec.supported_ops = [
#     tensorflow.lite.OpsSet.TFLITE_BUILTINS, # enable TensorFlow Lite ops.
#     tensorflow.lite.OpsSet.SELECT_TF_OPS # enable TensorFlow ops.
# ]

tflite_model = converter.convert()


with open('model.tflite', 'wb') as f:
    f.write(tflite_model)